package mmpathgen.callgraphanalysis;

import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Set;

import soot.G;
import soot.Scene;
import soot.SootClass;
import soot.SootMethod;
import soot.jimple.toolkits.annotation.purity.SootMethodFilter;
import soot.jimple.toolkits.callgraph.CallGraph;
import soot.jimple.toolkits.callgraph.Edge;
import soot.jimple.toolkits.callgraph.Sources;
import soot.jimple.toolkits.callgraph.Targets;
import soot.toolkits.graph.DirectedGraph;
import soot.util.HashMultiMap;
import soot.util.MultiMap;

public class CallGraphAnalysis {

	public static Targets getPossibleCallees(CallGraph cg, SootMethod src) {
		Targets tgts = new Targets(cg.edgesOutOf(src));

		return tgts;
	}

	public static List<SootMethod> getPossibleCallees(FilteredCallGraph fcg,
			SootMethod src) {
		List<SootMethod> tgts = fcg.getSuccsOf(src);

		return tgts;
	}

	public static Sources getPossibleCallers(CallGraph cg, SootMethod tgt) {
		Sources srcs = new Sources(cg.edgesInto(tgt));

		return srcs;
	}

	public static List<SootMethod> getPossibleCallers(FilteredCallGraph fcg,
			SootMethod tgt) {
		List<SootMethod> srcs = fcg.getPredsOf(tgt);

		return srcs;
	}

	@SuppressWarnings("unchecked")
	public static List<SootMethod> getTails(CallGraph cg) {
		Set<SootMethod> nodes;
		List<SootMethod> tails;

		List<SootMethod> heads = new LinkedList<SootMethod>();
		Iterator<SootClass> getClassesIt = Scene.v().getApplicationClasses()
				.iterator();
		while (getClassesIt.hasNext()) {
			SootClass appClass = getClassesIt.next();
			Iterator<SootMethod> getMethodsIt = appClass.getMethods()
					.iterator();
			while (getMethodsIt.hasNext()) {
				SootMethod method = getMethodsIt.next();
				if (method.getName().equals("main")) {
					heads.add(method);
					G.v().out.println("  |- will treat " + appClass.getName()
							+ "." + method.getName());
				}
			}
		}

		nodes = new HashSet<SootMethod>(heads);

		MultiMap s = new HashMultiMap();
		// MultiMap p = new HashMultiMap();

		Set<SootMethod> remain = new HashSet<SootMethod>(heads);
		int nb = 0;
		while (!remain.isEmpty()) {
			Set<SootMethod> newRemain = new HashSet<SootMethod>();
			Iterator<SootMethod> it = remain.iterator();
			while (it.hasNext()) {
				SootMethod m = it.next();
				Iterator<Edge> itt = cg.edgesOutOf(m);
				while (itt.hasNext()) {
					Edge edge = itt.next();
					SootMethod mm = edge.tgt();
					if (mm.isConcrete()) {
						if (nodes.add(mm))
							newRemain.add(mm);
						s.put(m, mm);
						// p.put(mm, m);
					}
				}
				nb++;
			}
			remain = newRemain;
		}

		tails = new LinkedList<SootMethod>();

		Iterator<SootMethod> it = nodes.iterator();
		while (it.hasNext()) {
			SootMethod x = it.next();
			Set<SootMethod> ss = s.get(x);
			if (ss.isEmpty())
				tails.add(x);
		}

		return tails;
	}

	public static void printCallGraphAsDot(DirectedGraph<SootMethod> cg,
			String mainClass, String cgaType, String filtType) {
		CallGraphToDotGraph cg2dg = new CallGraphToDotGraph();
		cg2dg.callGraphToDotGraph(cg, mainClass, cgaType, filtType);
	}

	public static void printCallGraphAsDot(DirectedGraph<SootMethod> cg,
			String mainClass, String cgaType, String filtType, String cfaType) {
		CallGraphToDotGraph cg2dg = new CallGraphToDotGraph();
		cg2dg.callGraphToDotGraph(cg, mainClass, cgaType, filtType, cfaType);
	}

	public static void printPossibleCallees(CallGraph cg, SootMethod src) {
		Targets tgts = new Targets(cg.edgesOutOf(src));
		while (tgts.hasNext()) {
			SootMethod tgt = (SootMethod) tgts.next();
			System.out.println(src + " may call" + tgt);
		}
	}

	public static void printPossibleCallees(FilteredCallGraph fcg, SootMethod src) {
		List<SootMethod> tgts = fcg.getSuccsOf(src);
		for (SootMethod tgt : tgts) {
			System.out.println(src + " may call" + tgt);
		}
	}

	public static void printPossibleCallers(CallGraph cg, SootMethod tgt) {
		Sources srcs = new Sources(cg.edgesInto(tgt));
		while (srcs.hasNext()) {
			SootMethod src = (SootMethod) srcs.next();
			System.out.println(tgt + " might be called by " + src);
		}
	}

	public static void printPossibleCallers(FilteredCallGraph fcg, SootMethod tgt) {
		List<SootMethod> srcs = fcg.getPredsOf(tgt);
		for (SootMethod src : srcs) {
			System.out.println(src + " may call" + tgt);
		}
	}

	public static FilteredCallGraph getFilteredCallGraph(CallGraph cg,
			SootMethodFilter filter) {
		// Find main methods
		List<SootMethod> heads = new LinkedList<SootMethod>();
		Iterator<SootClass> getClassesIt = Scene.v().getApplicationClasses()
				.iterator();
		while (getClassesIt.hasNext()) {
			SootClass appClass = getClassesIt.next();
			Iterator<SootMethod> getMethodsIt = appClass.getMethods()
					.iterator();
			while (getMethodsIt.hasNext()) {
				SootMethod method = getMethodsIt.next();
				if (method.getName().equals("main")) {
					heads.add(method);
					G.v().out.println("  |- will treat " + appClass.getName()
							+ "." + method.getName());
				}
			}
		}

		FilteredCallGraph fcg = new FilteredCallGraph(cg, filter, heads.iterator());
		return fcg;
	}
	
	public static BoundedCallGraph getBoundedCallGraph(CallGraph cg,
			int mtdBound) {
		// Find main methods
		List<SootMethod> heads = new LinkedList<SootMethod>();
		Iterator<SootClass> getClassesIt = Scene.v().getApplicationClasses()
				.iterator();
		while (getClassesIt.hasNext()) {
			SootClass appClass = getClassesIt.next();
			Iterator<SootMethod> getMethodsIt = appClass.getMethods()
					.iterator();
			while (getMethodsIt.hasNext()) {
				SootMethod method = getMethodsIt.next();
				if (method.getName().equals("main")) {
					heads.add(method);
					G.v().out.println("  |- will treat " + appClass.getName()
							+ "." + method.getName());
				}
			}
		}

		BoundedCallGraph bcg = new BoundedCallGraph(cg, heads.iterator(), mtdBound);
		return bcg;
	}

	public static UnfilteredCallGraph getUnfilteredCallGraph(CallGraph cg) {
		// Find main methods
		List<SootMethod> heads = new LinkedList<SootMethod>();
		Iterator<SootClass> getClassesIt = Scene.v().getApplicationClasses()
				.iterator();
		while (getClassesIt.hasNext()) {
			SootClass appClass = getClassesIt.next();
			Iterator<SootMethod> getMethodsIt = appClass.getMethods()
					.iterator();
			while (getMethodsIt.hasNext()) {
				SootMethod method = getMethodsIt.next();
				if (method.getName().equals("main")) {
					heads.add(method);
					G.v().out.println("  |- will treat " + appClass.getName()
							+ "." + method.getName());
				}
			}
		}
		UnfilteredCallGraph ucg = new UnfilteredCallGraph(cg, heads.iterator());
		return ucg;
	}

}